import datetime
import numpy as np
import sys
import os
import re
import cv2 as cv
import matplotlib.pyplot as plt


def sep(label = '', cnt=32):
    print('-' * cnt, label, '-' * cnt, sep='')


def print_numpy_ndarray_info(arr, name):
    print(f'min: {arr.min():.2e},\tmax: {arr.max():.2e},\tmean: {arr.mean():.2e},\tstd: {arr.std():.2e},\tdtype: {arr.dtype},\tshape: {arr.shape},\t{name}')


def print_numpy_ndarray_info_simple(arr, name):
    print(f'{name}: dtype: {arr.dtype}, shape: {arr.shape}')


def check_shape(var, name, file=sys.stderr):
    print(f'{type(var)} {name}.shape = {np.shape(var)}', file=file)


def check_shape_stdout(var, name):
    check_shape(var, name, sys.stdout)


def flush():
    sys.stdout.flush()
    sys.stderr.flush()


def line_2_dots(x1, y1, x2, y2, x):
    return (x - x1) * (y2 - y1) / (x2 - x1) + y1


def select_in_range(base, decrease, increase, low_border, up_border):
    if decrease < 0:
        raise Exception('Invalid decrease amount.')
    if increase < 0:
        raise Exception('Invalid increase amount.')
    if low_border > up_border:
        raise Exception('Invalid borders')
    if decrease + increase > up_border - low_border:
        raise Exception("Out of range.")
    a = base - decrease
    b = base + increase
    if low_border <= a and b <= up_border:
        return (a, b), None
    if a < low_border:
        return (low_border, b), (up_border - (low_border - a), up_border)
    if up_border < b:
        return (low_border, low_border + (b - up_border)), (a, up_border)


def select_color(base, decrease, increase):
    return select_in_range(base, decrease, increase, 0, 180)


def imzoom_rect(img, tgt_rect, inter=cv.INTER_CUBIC):
    tgt_w, tgt_h = tgt_rect
    tgt_h, tgt_w = int(tgt_h), int(tgt_w)
    ori_h, ori_w = img.shape[:2]
    ori_ch_slice = img.shape[2:3]
    tgt_rate = tgt_w / tgt_h
    ori_rate = ori_w / ori_h
    if ori_rate < tgt_rate:
        adj_h = ori_h
        adj_w = adj_h * tgt_rate
        padding = abs(adj_w - ori_w)
        padding01, padding02 = divide_int(padding)
        img = np.concatenate([
            np.zeros((adj_h, padding01, *ori_ch_slice), dtype=img.dtype),
            img,
            np.zeros((adj_h, padding02, *ori_ch_slice), dtype=img.dtype),
        ], axis=1)
    else:
        adj_w = ori_w
        adj_h = adj_w / tgt_rate
        padding = abs(adj_h - ori_h)
        padding01, padding02 = divide_int(padding)
        img = np.concatenate([
            np.zeros((padding01, adj_w, *ori_ch_slice), dtype=img.dtype),
            img,
            np.zeros((padding02, adj_w, *ori_ch_slice), dtype=img.dtype),
        ], axis=0)

    img = imzoom(img, tgt_h, h_or_w='h', inter=inter)
    return img


def imzoom(img, tgt_len, h_or_w='h', inter=cv.INTER_LINEAR):
    shape_tuple = img.shape
    h_ori = shape_tuple[0]
    w_ori = shape_tuple[1]
    if 'h' == h_or_w:
        rate = tgt_len / h_ori
    else:
        rate = tgt_len / w_ori
    img = cv.resize(img, None, fx=rate, fy=rate, interpolation=inter)
    return img


# （16）用函数cal_convoluation_size（）计算卷积尺寸
def cal_convolution_size(size, kernel_size):
    return size + 1 - kernel_size


def no_padding(img, ksize):
    img_shape = np.shape(img)
    H, W = img_shape[0], img_shape[1]
    HK, WK = ksize[0], ksize[1]
    PH = (HK - 1) // 2
    PW = (WK - 1) // 2
    return img[PH:H - PH, PW:W - PW]


# （17）用函数convoluation(image, kernel)计算卷积结果
def convolution(image, kernel):
    shape = np.shape(image)
    h = shape[0]
    w = shape[1]
    ch = None
    if len(shape) > 2:
        ch = shape[2]
    k_shape = np.shape(kernel)
    kh = k_shape[0]
    kw = k_shape[1]
    res_h = cal_convolution_size(h, kh)
    res_w = cal_convolution_size(w, kw)
    if not ch:
        res_mat = np.zeros((res_h, res_w), dtype=np.float32)
    else:
        res_mat = np.zeros((res_h, res_w, ch), dtype=np.float32)
    for i in range(res_h):
        for j in range(res_w):
            for ii in range(kh):
                for jj in range(kw):
                    if not ch:
                        res_mat[i][j] += image[ii + i][jj + j] * kernel[ii][jj]
                    else:
                        for kk in range(ch):
                            res_mat[i][j][kk] += image[ii + i][jj + j][kk] * kernel[ii][jj]
    return res_mat


def convolution_numpy(image, kernel):
    shape = np.shape(image)
    h = shape[0]
    w = shape[1]
    ch = None
    if len(shape) > 2:
        ch = shape[2]
    k_shape = np.shape(kernel)
    kh = k_shape[0]
    kw = k_shape[1]
    res_h = cal_convolution_size(h, kh)
    res_w = cal_convolution_size(w, kw)
    if not ch:
        res_mat = np.zeros((res_h, res_w), dtype=np.float32)
    else:
        res_mat = np.zeros((res_h, res_w, ch), dtype=np.float32)
    for i in range(res_h):
        for j in range(res_w):
            for ii in range(kh):
                for jj in range(kw):
                    if not ch:
                        res_mat[i, j] += image[ii + i, jj + j] * kernel[ii, jj]
                    else:
                        for kk in range(ch):
                            res_mat[i, j, kk] += image[ii + i, jj + j, kk] * kernel[ii, jj]
    return res_mat


def convolution_numpy_acc(image, kernel):
    shape = np.shape(image)
    h = shape[0]
    w = shape[1]
    ch = None
    if len(shape) > 2:
        ch = shape[2]
    k_shape = np.shape(kernel)
    kh = k_shape[0]
    kw = k_shape[1]
    res_h = cal_convolution_size(h, kh)
    res_w = cal_convolution_size(w, kw)
    if not ch:
        res_mat = np.zeros((res_h, res_w), dtype=np.float32)
    else:
        res_mat = np.zeros((res_h, res_w, ch), dtype=np.float32)
    for i in range(res_h):
        for j in range(res_w):
            if not ch:
                res_mat[i, j] = (image[i:kh + i, j:kw + j] * kernel).sum()
            else:
                for kk in range(ch):
                    res_mat[i, j, kk] = (image[i:kh + i, j:kw + j, kk] * kernel).sum()
    return res_mat


def median_blur(image, k_shape):
    shape = image.shape
    h = shape[0]
    w = shape[1]
    ch = None
    if len(shape) > 2:
        ch = shape[2]
    kh = k_shape[0]
    kw = k_shape[1]
    res_h = cal_convolution_size(h, kh)
    res_w = cal_convolution_size(w, kw)
    if not ch:
        res_mat = np.zeros((res_h, res_w), dtype=np.float32)
    else:
        res_mat = np.zeros((res_h, res_w, ch), dtype=np.float32)
    for i in range(res_h):
        for j in range(res_w):
            if not ch:
                res = np.median(image[i:i + kh, j:j + kw])
            else:
                res = np.zeros((ch,))
                for kk in range(ch):
                    res[kk] = np.median(image[i:i + kh, j:j + kw, kk])
            if not ch:
                res_mat[i, j] = res
            else:
                res_mat[i, j, :] = res
    return res_mat


def make_salt_and_pepper(img_, SNR):
    img = img_.copy()
    mask = np.random.choice((0, 1, 2), img.shape, True, (SNR, (1 - SNR) / 2, (1 - SNR) / 2))  # ATTENTION API np.random.choice
    img[mask == 1] = 0
    img[mask == 2] = 255
    return img


def make_salt_and_pepper_black_white(img_, SNR):
    img = img_.copy()
    mask = np.random.choice((0, 1, 2), (img.shape[0], img.shape[1]), True, (SNR, (1 - SNR) / 2, (1 - SNR) / 2))
    img[mask == 1] = 0
    img[mask == 2] = 255
    return img


def my_erode_or_dilate(img, kenel_size, iterations=1, is_dilate=False):
    for i in range(iterations):
        img = my_erode_or_dilate_once(img, kenel_size, is_dilate=is_dilate)
    return img


def my_erode_or_dilate_once(img, kenel_size, is_dilate=False):
    shape = img.shape
    h = shape[0]
    w = shape[1]
    ch = None
    if len(shape) > 2:
        ch = shape[2]
    kh, kw = kenel_size

    ph = (kh - 1) // 2
    pw = (kw - 1) // 2

    # padding
    if ch is None:
        img = np.concatenate((np.zeros((ph, w), dtype=np.uint8),
                              img,
                              np.zeros((ph, w), dtype=np.uint8)), axis=0)
        img = np.concatenate((np.zeros((ph * 2 + h, pw), dtype=np.uint8),
                              img,
                              np.zeros((ph * 2 + h, pw), dtype=np.uint8)), axis=1)
    else:
        img = np.concatenate((np.zeros((ph, w, ch), dtype=np.uint8),
                              img,
                              np.zeros((ph, w, ch), dtype=np.uint8)), axis=0)
        img = np.concatenate((np.zeros((ph * 2 + h, pw, ch), dtype=np.uint8),
                              img,
                              np.zeros((ph * 2 + h, pw, ch), dtype=np.uint8)), axis=1)

    if is_dilate:
        init_v = 0
    else:
        init_v = 255
    if ch is None:
        res = np.full((h, w), init_v, np.uint8)
    else:
        res = np.full((h, w, ch), init_v, np.uint8)

    for i in range(h):
        for j in range(w):
            for ii in range(kh):
                for jj in range(kw):
                    if ch is None:
                        if is_dilate:
                            res[i, j] |= img[i + ii, j + jj]
                        else:
                            res[i, j] &= img[i + ii, j + jj]
                    else:
                        for kk in range(ch):
                            if is_dilate:
                                res[i, j, kk] |= img[i + ii, j + jj, kk]
                            else:
                                res[i, j, kk] &= img[i + ii, j + jj, kk]

    return res


def inverse_transform_weights(w, b, mu, sigma):
    w = w.ravel()
    b = b.ravel()[0]
    mu = mu.ravel()
    sigma = sigma.ravel()
    print('w:', w, 'b:', b)

    mu_x = mu[:-1]
    sigma_x = sigma[:-1]
    mu_y = mu[-1]
    sigma_y = sigma[-1]
    # print('mu_x', mu_x)
    # print('sigma_x', sigma_x)
    # print('mu_y', mu_y)
    # print('sigma_y', sigma_y)

    w_ori = w * sigma_y / sigma_x

    xsum = (w * mu_x / sigma_x).sum()
    # print('xsum', xsum)
    b_ori = (b - xsum) * sigma_y + mu_y

    return w_ori, b_ori


repr2valid_name_regexp = re.compile(r'[^0-9a-zA-Z_]')
def repr2valid_name(var):
    repr = str(var)
    name = repr2valid_name_regexp.sub('_', repr)
    return name


def split_img(img, h_num, w_num):
    h, w = img.shape[:2]
    h_unit = h / h_num
    w_unit = w / w_num
    for i in range(1, h_num):
        cv.line(img, (0, int(h_unit * i)), (w, int(h_unit * i)), rand_color(), 1)
    for i in range(1, w_num):
        cv.line(img, (int(w_unit * i), 0), (int(w_unit * i), h), rand_color(), 1)
    return img


def xywh2pts(x, y, w, h, idx_grids, w_input_size=448, h_input_size=448, w_grids=7, h_grids=7):
    gx = idx_grids % w_grids
    gy = idx_grids // w_grids
    w_grid_size = w_input_size / w_grids
    h_grid_size = h_input_size / h_grids

    base_x = gx * w_grid_size
    base_y = gy * h_grid_size
    x = base_x + w_grid_size * x
    y = base_y + h_grid_size * y
    w *= w_input_size
    h *= h_input_size

    x1 = x - w / 2
    y1 = y - h / 2
    x2 = x1 + w
    y2 = y1 + h

    x1 = int(x1)
    y1 = int(y1)
    x2 = int(x2)
    y2 = int(y2)
    return (x1, y1), (x2, y2)


def show_cv_img_by_plt(plt_info, img, title="no title", trans=None, **kwargs):
    plt_info['spn'] += 1
    plt.subplot(plt_info['spr'], plt_info['spc'], plt_info['spn'])
    if trans is not None:
        img = trans(img)
    plt.imshow(img, **kwargs)
    plt.axis('off')
    plt.title(title)
